package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.common;

import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1.MnistTrain;

import java.io.*;
import java.util.zip.GZIPInputStream;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;

public class MNISTDataExtractor {

    public static void main(String[] args) {
        String imagesFilePath = MnistTrain.BASE_PATH + "/mnist_png/training/train-images-idx3-ubyte.gz";
        String labelsFilePath = MnistTrain.BASE_PATH + "/mnist_png/training/train-labels-idx1-ubyte.gz";
        String outputDir = MnistTrain.BASE_PATH + "/mnist_png/decompression";

        // 解压图像数据集
        extractImages(imagesFilePath, outputDir);

        // 解压标签数据集
        extractLabels(labelsFilePath, outputDir);

        System.out.println("MNIST data extraction completed.");
    }

    private static void extractImages(String filePath, String outputDir) {
        try {
            FileInputStream fileInputStream = new FileInputStream(filePath);
            GZIPInputStream gzipInputStream = new GZIPInputStream(fileInputStream);
            DataInputStream dataInputStream = new DataInputStream(gzipInputStream);

            // 读取文件头
            int magicNumber = dataInputStream.readInt();
            int numImages = dataInputStream.readInt();
            int numRows = dataInputStream.readInt();
            int numCols = dataInputStream.readInt();

            // 逐个解压图像并保存为文件
            for (int i = 0; i < numImages; i++) {
                byte[] imageBuffer = new byte[numRows * numCols];
                dataInputStream.readFully(imageBuffer);

                // 创建图像对象
                BufferedImage image = new BufferedImage(numCols, numRows, BufferedImage.TYPE_BYTE_GRAY);

                // 设置图像像素值
                for (int row = 0; row < numRows; row++) {
                    for (int col = 0; col < numCols; col++) {
                        int pixelValue = imageBuffer[row * numCols + col] & 0xFF;
                        image.setRGB(col, row, pixelValue | pixelValue << 8 | pixelValue << 16);
                    }
                }

                // 创建文件名
                String fileName = outputDir + "/" + i + ".png";

                // 保存图像到文件
                ImageIO.write(image, "png", new File(fileName));
            }

            dataInputStream.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private static void extractLabels(String filePath, String outputDir) {
        try {
            FileInputStream fileInputStream = new FileInputStream(filePath);
            GZIPInputStream gzipInputStream = new GZIPInputStream(fileInputStream);
            DataInputStream dataInputStream = new DataInputStream(gzipInputStream);

            // 读取文件头
            int magicNumber = dataInputStream.readInt();
            int numLabels = dataInputStream.readInt();

            // 逐个解压标签并保存为文件
            for (int i = 0; i < numLabels; i++) {
                byte label = dataInputStream.readByte();

                // 创建文件名
                String fileName = outputDir + "/" + i + ".txt";

                // 创建标签文件并保存标签
                FileOutputStream fileOutputStream = new FileOutputStream(fileName);
                BufferedOutputStream bufferedOutputStream = new BufferedOutputStream(fileOutputStream);
                bufferedOutputStream.write(String.valueOf(label).getBytes());
                bufferedOutputStream.close();
            }

            dataInputStream.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}